#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <getopt.h>
#include <iptables.h>
#include <string.h>

#include "ipt_ipcontext.h"

int parseArguments(int c, char **argv, unsigned int *flags, struct ipt_entry_match **match);
void final_check_Args(unsigned int flags);

static struct option opts[] = {
	{.name = "ipcontext_port"       , .has_arg = 1, .flag = 0, .val = '1'},
	{.name = "ipcontext_protocol"   , .has_arg = 1, .flag = 0, .val = '2'},
	{.name = "ipcontext_log"        , .has_arg = 0, .flag = 0, .val = '3'},
	{.name = "ipcontext_regexp"     , .has_arg = 1, .flag = 0, .val = '4'},	
	{.name = "ipcontext_dir" 	, .has_arg = 1, .flag = 0, .val = '5'},
}; 

#define CAST_MATCH(match_info) ((struct ipt_context_info*)(match_info->data))

/* Checking that match_info's port is 'desiredPort' */
void testPort(struct ipt_entry_match* match_info, int desiredPort)
{
	if (CAST_MATCH(match_info)->port != desiredPort) {
		printf("error: port wasn't set to %d\n", desiredPort);
	}
	else {
		printf("port is set to %d\n", desiredPort);
	}
}

/* Checking that match_info's protocol is 'desiredProtocol' */ 
void testProtocol(struct ipt_entry_match* match_info, char* desiredProtocol)
{
	if (strcmp(CAST_MATCH(match_info)->protocol, desiredProtocol) != 0) {
		printf("error: protocol wasn't set to %s\n", desiredProtocol);
	}
	else {
		printf("protocol is set to %s\n", desiredProtocol);
	} 
}

/* Checking that match_info's logger's state is 'desiredState' */
void testLoggerState(struct ipt_entry_match* match_info, int desiredState)
{
	if (CAST_MATCH(match_info)->flags & IPCONTEXT_LOG) {
		if (desiredState == 1) {
			printf("logger is ON\n"); 
		}
		else {
			printf("error: logger is ON\n"); 
		}
	}
	else {
		if (desiredState == 0) {
			printf("logger is OFF\n"); 
		}
		else {
			printf("error: logger is OFF\n"); 
		}
	} 
}

/* Checking that match_info's regular expression is 'desiredRegexp' */
void testRegularExpression(struct ipt_entry_match* match_info, char* desiredRegexp)
{
	if (strcmp(CAST_MATCH(match_info)->regular_exp, desiredRegexp) != 0) {
		printf("regexp error: didn't manage to set %s\n", desiredRegexp);
	}
	else {
		printf("regexp is set to %s\n", desiredRegexp);
	} 
}

/* Checking that match_info's direction is 'desiredDirection' */
void testDirection(struct ipt_entry_match* match_info, int desiredDirection)
{
	if (CAST_MATCH(match_info)->flags & IPCONTEXT_IN) {
		if (desiredDirection == 1) {
			printf("directon is set to IN\n");
		}
		else {
			printf("error: direction is set to IN\n");
		}
	}
	else {
		if (desiredDirection == 0) {
			printf("directon is set to OUT\n");
		}
		else {
			printf("error: direction is set to OUT\n");
		}
	} 
}

void testFinalCheck(unsigned int flags)
{
	printf("\nRunning final check test:\n");
	printf("========================\n");
	if ((flags & IPCONTEXT_PORT) 
		&& (flags & IPCONTEXT_PROTOCOL)
		&& (flags & IPCONTEXT_REGEXP)) {
		printf("final check result: test is successful\n");	
		return;	
	}
	printf("Invalid Params! %s %s %s ",
	      ( flags & IPCONTEXT_PORT )    ? " " : "\n --ipcontext_port   <port> is missing" ,
	      ( flags & IPCONTEXT_REGEXP )  ? " " : "\n --ipcontext_regexp <regexp>is missing" ,
	      ( flags & IPCONTEXT_PROTOCOL )? " " : "\n --ipcontext_protocol <protocol> is missing" );
}

/* Checking that regexps longer than REGEXP_SIZE are not allowed */
int test0() 
{
	char* long_regexp = (char*)malloc(REGEXP_SIZE + 1);
	if (long_regexp == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	memset(long_regexp, 'a', REGEXP_SIZE);
	long_regexp[REGEXP_SIZE] = '\0';

	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;
	int c;
	unsigned int flags = 0;
	int option_index = 0;	

	char* argv[] = {"dummy","--ipcontext_regexp",long_regexp};
	printf("The test: --ipcontext_regexp <long regexp>\n");
	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");
	while ((c = getopt_long(3, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {

		int rc = parseArguments(c, argv, &flags, &match_info);
		if (c == '0' + 4) {
			if (rc != 0) {
				printf("Error rejecting long regular expression!\n");
			}
		}
	}
	free(long_regexp);
}

/* Checking that 'parse' manages to set port to 80 */
int test1()
{
	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;
	int c;
	unsigned int flags = 0;
	int option_index = 0;	

	char* argv[] = {"dummy","--ipcontext_port","80"};

	printf("The test: --ipcontext_port 80\n");
	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");
	while ((c = getopt_long(3, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {
		parseArguments(c, argv, &flags, &match_info);
	}
	
	testPort(match_info, 80);
	testLoggerState(match_info, 0);
	final_check_Args(flags);
	testFinalCheck(flags);
	free(match);
}

/* Checking that 'parse' manages to set protocol to http */
int test2()
{
	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;
	int c;
	unsigned int flags = 0;
	int option_index = 0;	

	char* argv[] = {"dummy","--ipcontext_protocol","http"};
	printf("The test: --ipcontext_protocol http\n");
	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");
	while ((c = getopt_long(3, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {
		parseArguments(c, argv, &flags, &match_info);
	}

	testProtocol(match_info, "http");
	testLoggerState(match_info, 0);
	final_check_Args(flags);
	testFinalCheck(flags);
	free(match);
}

/* Checking that 'parse' manages to set protocol's regular expression to 'expression' */
int test3()
{
	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;
	int c;
	unsigned int flags = 0;
	int option_index = 0;	

	char* argv[] = {"dummy","--ipcontext_regexp","REGEXP"};
	printf("The test: --ipcontext_regexp REGEXP\n");
	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");
	while ((c = getopt_long(3, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {
		parseArguments(c, argv, &flags, &match_info);
	}
	testRegularExpression(match_info, "REGEXP");
	testLoggerState(match_info, 0);
	final_check_Args(flags);
	testFinalCheck(flags);
	free(match);
}

/* Checking that 'parse' manages to set logger's state to ON */
int test4()
{
	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;
	int c;
	unsigned int flags = 0;
	int option_index = 0;	

	char* argv[] = {"dummy","--ipcontext_log"};
	printf("The test: --ipcontext_log\n");
	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");
	while ((c = getopt_long(2, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {
		parseArguments(c, argv, &flags, &match_info);
	}
	testLoggerState(match_info, 1);
	final_check_Args(flags);
	testFinalCheck(flags);
	free(match);
}

/* Checking that 'parse' manages to set rule's direction to IN */
int test5()
{
	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;
	int c;
	unsigned int flags = 0;
	int option_index = 0;	

	char* argv[] = {"dummy","--ipcontext_dir","1"};
	printf("The test: --ipcontext_direction 1\n");
	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");
	while ((c = getopt_long(3, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {
		parseArguments(c, argv, &flags, &match_info);
	}
	testDirection(match_info, 1);
	testLoggerState(match_info, 0);
	final_check_Args(flags);
	testFinalCheck(flags);
	free(match);
}

int test6()
{
	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;
	int c;
	unsigned int flags = 0;
	int option_index = 0;	

	char* argv[] = {"dummy","--ipcontext_port","80","--ipcontext_protocol" ,
			"http", "--ipcontext_regexp","REGEXP"};
	printf("The test:--ipcontext_port 80 --ipcontext_protocol http --ipcontext_regexp REGEXP\n");
	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");
	while ((c = getopt_long(7, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {
		parseArguments(c, argv, &flags, &match_info);
	}
	testPort(match_info, 80);
	testProtocol(match_info, "http");
	testRegularExpression(match_info, "REGEXP");
	testLoggerState(match_info, 0);
	final_check_Args(flags);
	testFinalCheck(flags);
	free(match);
} 

int test7()
{
	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;
	int c;
	unsigned int flags = 0;
	int option_index = 0;	

	char* argv[] = {"dummy","--ipcontext_port","80","--ipcontext_protocol" ,
			"http", "--ipcontext_regexp","REGEXP","--ipcontext_log"};
	printf("The test:--ipcontext_port 80 --ipcontext_protocol http --ipcontext_regexp REGEXP --ipcontext_log\n");
	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");
	while ((c = getopt_long(8, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {
		parseArguments(c, argv, &flags, &match_info);
	}
	testPort(match_info, 80);
	testProtocol(match_info, "http");
	testRegularExpression(match_info, "REGEXP");
	testLoggerState(match_info, 1);
	final_check_Args(flags);
	testFinalCheck(flags);
	free(match);
} 

int test8()
{
	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;
	int c;
	unsigned int flags = 0;
	int option_index = 0;	

	char* argv[] = {"dummy","--ipcontext_port","80","--ipcontext_protocol" ,
			"http", "--ipcontext_regexp","REGEXP","--ipcontext_dir","1"};
	printf("The test:--ipcontext_port 80 --ipcontext_protocol http --ipcontext_regexp REGEXP --ipcontext_dir 1\n");
	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");
	while ((c = getopt_long(9, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {
		parseArguments(c, argv, &flags, &match_info);
	}
	testPort(match_info, 80);
	testProtocol(match_info, "http");
	testRegularExpression(match_info, "REGEXP");
	testLoggerState(match_info, 0);
	testDirection(match_info, 1);
	final_check_Args(flags);
	testFinalCheck(flags);
	free(match);
} 

int test9()
{
	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;
	int c;
	unsigned int flags = 0;
	int option_index = 0;	

	char* argv[] = {"dummy","--ipcontext_port","80","--ipcontext_protocol" ,
			"http", "--ipcontext_regexp","REGEXP","--ipcontext_log","--ipcontext_dir","1"};
	printf("The test:--ipcontext_port 80 --ipcontext_protocol http --ipcontext_regexp REGEXP --ipcontext_log --ipcontext_dir 1\n");
	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");
	while ((c = getopt_long(10, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {
		parseArguments(c, argv, &flags, &match_info);
	}
	testPort(match_info, 80);
	testProtocol(match_info, "http");
	testRegularExpression(match_info, "REGEXP");
	testLoggerState(match_info, 1);
	testDirection(match_info, 1);
	final_check_Args(flags);
	testFinalCheck(flags);
	free(match);
} 

/* may be risky!! */
char* generateRandomString(int minlen, int maxlen)
{
	int i, r;
	int len = minlen + (rand() % (maxlen - minlen + 1));
	char* result = (char*)malloc(sizeof(char) * (len + 1));
	if (result == NULL) {
		return NULL;
	}
	result[len] = '\0';
	for(i = 0; i < len; i++) {
		r = (rand() % 3);
		if (r == 0) { // generate a digit (0-9) 
			result[i] = '0' + (rand() % 10);
		}
		else if (r == 1) { // lower letter 
			result[i] = 'a' + (rand() % 26);
		}
		else { // capital leter 
			result[i] = 'A' + (rand() % 26);
		}	
	} 
	return result;
}
					
char* allocateString(int len)
{
	char* str = (char*)malloc(sizeof(char) * (len + 1));
	if (str == NULL) {
		return NULL;
	}
	else {
		str[len] = '\0';
		return str;
	}
}

int digits(int num)
{
	int cnt = 0;
	while (num != 0) {
		num /= 10;
		cnt++;
	}
	return cnt;
}

char* intToString(int num)
{
	int i, n = digits(num);
	char* numStr = allocateString(n);
	numStr[n - 1] = '\0';
	for(i = 0; i < n; i++) {
		numStr[n - 1 - i] = '0' + num % 10;
		num /= 10;
	}
	return numStr;
}

/*
void test10()
{
	srand((unsigned)time(NULL));
	int port = (rand() % 65535) + 1; // generate a port number
	char* protocol   = generateRandomString(3, 50); // generate a protocol 
	char* regexp     = generateRandomString(10, 200); // generate regexp 
	int loggerON     = rand() % 2; // generate logger state 
	int directionIN  = rand() % 2; // generate direction 
	char* argv[10]; 
	int argc = 7;
	struct ipt_context_info* match = (struct ipt_context_info*)malloc(sizeof(struct ipt_context_info));
	if (match == NULL) {
		fprintf(stderr, "Memory allocation failed!\n");
		exit(1);
	}
	struct ipt_entry_match* match_info = (struct ipt_entry_match*)match;

	int c, i;
	unsigned int flags = 0;
	int option_index = 0;	

	argv[0] = allocateString(strlen("dummy"));
	strcpy(argv[0], "dummy"); 
	argv[1] = allocateString(strlen("--ipcontext_port"));
	strcpy(argv[1],"--ipcontext_port");
	argv[2] = allocateString(strlen("65535"));
	char* portStr = intToString(port);
	strcpy(argv[2], portStr);
	argv[3] = allocateString(strlen("--ipcontext_protocol"));
	strcpy(argv[3], "--ipcontext_protocol");
	argv[4] = allocateString(strlen(protocol));
	strcpy(argv[4], protocol);
	argv[5] = allocateString(strlen("--ipcontext_regexp"));
	strcpy(argv[5], "--ipcontext_regexp");
	argv[6] = allocateString(strlen(regexp));
	strcpy(argv[6], regexp); 

	if (loggerON) {	
		argc++;
		argv[7] = allocateString(strlen("--ipcontext_log"));
		strcpy(argv[7], "--ipcontext_log");
	}
	else {
		argv[7] = allocateString(1);
		strcpy(argv[7], "");
	}
	if (directionIN) { 
		argc+=2;
		argv[8] = allocateString(strlen("--ipcontext_dir"));
		strcpy(argv[8], "--ipcontext_dir");
		argv[9] = allocateString(strlen("1"));
		strcpy(argv[9], "1");
	}

	printf("\nThe test: --ipcontext_port %d --ipcontext_protocol %s --ipcontext_regexp %s", port, protocol, regexp);
	printf((loggerON ? " --ipcontext_log" : ""));
	printf((directionIN ? " --ipcontext_dir 1\n" : "\n"));

	printf("Running ipcontext parse function test:\n");
	printf("=====================================\n");

	while ((c = getopt_long(argc, argv,
		"ipcontext_port:ipcontext_protocol:ipcontext_log:ipcontext_regexp:ipcontext_dir",
		opts, &option_index)) != -1) {
		parseArguments(c, argv, &flags, &match_info);
	}
	testPort(match_info, port);
	testProtocol(match_info, protocol);
	testRegularExpression(match_info, regexp);
	if (loggerON) {
		testLoggerState(match_info, 1);
	}
	else {
		testLoggerState(match_info, 0);
	}

	if (directionIN) {
		testDirection(match_info, 1);
 	}
	else {
		testDirection(match_info, 0);
	} 
	final_check_Args(flags);
	testFinalCheck(flags);
	free(portStr);
	free(protocol);
	free(regexp);
	free(match);
}  */

int main(int argc, char** argv)
{
	printf("\nTest %s:\n", argv[1]);
	printf("======\n");
	switch (atoi(argv[1])) {
		case 0:	test0();
			break;
		case 1:	test1();
			break;
		case 2: test2();
			break;	
		case 3: test3();
			break;
		case 4: test4();
			break;
		case 5: test5();
			break;
		case 6: test6();
			break;
		case 7: test7();
			break;
		case 8: test8();
			break;
		case 9: test9();
			break;
		//case 10: test10(); // random test
		//	break;
		default:
			return 0;
	}
	printf("\nTest %s is done.\n", argv[1]);
	return 1;
}
